import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
from torch.autograd import Variable
from torch.nn import functional as F
from torchvision import datasets, transforms
# from torchvision.utils import save_image

BATCH_SIZE = 64
EPOCH = 10


class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()
        self.fc1 = nn.Linear(784, 400)
        self.fc21 = nn.Linear(400, 20)
        self.fc22 = nn.Linear(400, 20)
        self.fc3 = nn.Linear(20, 400)
        self.fc4 = nn.Linear(400, 784)

        self.relu = nn.ReLU()
        self.sigmoid = nn.Sigmoid()

    def encode(self, x):
        h1 = self.relu(self.fc1(x))
        return self.fc21(h1), self.fc22(h1)

    # todo figure out why need reparameterize?
    def reparameterize(self, mu, logvar):
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = Variable(std.data.new(std.size()).normal_())
            return eps.mul(std).add_(mu)
        else:
            return mu

    def decode(self, z):
        h3 = self.relu(self.fc3(z))
        return self.sigmoid(self.fc4(h3))

    def forward(self, x):
        mu, logvar = self.encode(x.view(-1, 784))
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar


def loss_function(recon_x, x, mu, logvar):
    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784))
    # KL divergence
    # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD = -0.5 * torch.sum(1.0 + logvar - mu.pow(2) - logvar.exp())
    KLD /= BATCH_SIZE * 784
    return BCE + KLD


if __name__ == '__main__':
    import sys

    torch.cuda.manual_seed(10)
    data_dir = sys.argv[1]
    model_file = sys.argv[2]

    train_data = datasets.MNIST(root=data_dir, train=True,
                                transform=transforms.ToTensor(),
                                download=False)

    train_loader = data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE,
                                   shuffle=True)

    # test_data = datasets.MNIST(root=data_dir, train=False,
    #                            transform=transforms.ToTensor(),
    #                            download=False)
    #
    # test_loader = data.DataLoader(dataset=test_data, batch_size=BATCH_SIZE,
    #                               shuffle=True)

    model = VAE().cuda()

    optimizer = optim.Adam(model.parameters(), lr=1e-3)

    all_batches = len(train_loader)

    for e in range(EPOCH):
        model.train()
        train_loss = 0

        for batch_idx, (data, _) in enumerate(train_loader):
            data = Variable(data).cuda()

            optimizer.zero_grad()

            recon_batch, mu, logvar = model(data)

            loss = loss_function(recon_batch, data, mu, logvar)
            loss.backward()

            train_loss += loss.data[0]
            optimizer.step()

            print("epoch: {} step: {}/{} loss: {:.2f}".format(e, batch_idx, all_batches, loss.data[0]))

    torch.save(model.state_dict(), model_file)

    # todo added test for encode and decode
